import os
import random
import matplotlib.pyplot as plt
import torch
import torchvision
from torchvision import transforms
from data.dataset import load_cifar10
from models.resnet import ResNet18
from utils.config import config

def denormalize(tensor, mean=(0.4914, 0.4822, 0.4465), std=(0.2023, 0.1994, 0.2010)):
    """反归一化图像，用于显示"""
    for t, m, s in zip(tensor, mean, std):
        t.mul_(s).add_(m)
    return tensor

def visualize_predictions(model, val_loader, classes, num_images=8, save_path=None):
    """
    可视化模型预测结果
    
    参数:
        model: 训练好的模型
        val_loader: 验证集数据加载器
        classes: 类别名称列表
        num_images: 要可视化的图像数量
        save_path: 保存可视化结果的路径(可选)
    """
    # 获取一批验证数据
    dataiter = iter(val_loader)
    images, labels = next(dataiter)
    
    # 随机选择num_images张图片
    indices = random.sample(range(len(images)), num_images)
    images = images[indices]
    labels = labels[indices]
    
    # 将图像移动到设备上并进行预测
    images = images.to(config.DEVICE)
    with torch.no_grad():
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)
    
    # 将图像移回CPU并反归一化
    images = images.cpu()
    images = [denormalize(img) for img in images]
    
    # 创建图像网格
    fig, axes = plt.subplots(nrows=2, ncols=num_images//2, figsize=(20, 8))
    if num_images % 2 != 0:
        fig, axes = plt.subplots(nrows=2, ncols=(num_images//2)+1, figsize=(20, 8))
    
    # 调整子图间距
    plt.subplots_adjust(wspace=0.5, hspace=0.5)
    
    # 绘制每张图像及其预测结果
    for idx, ax in enumerate(axes.flat):
        if idx >= num_images:
            break
            
        # 显示图像
        img = images[idx].permute(1, 2, 0).numpy()
        ax.imshow(img)
        
        # 设置标题
        true_label = classes[labels[idx]]
        pred_label = classes[predicted[idx]]
        title_color = 'green' if true_label == pred_label else 'red'
        
        ax.set_title(f"True: {true_label}\nPred: {pred_label}", color=title_color)
        ax.axis('off')
    
    # 添加整体标题
    fig.suptitle('Model Predictions on Validation Set\n(Green=Correct, Red=Incorrect)', 
                 fontsize=16, y=1.05)
    
    # 保存或显示图像
    if save_path:
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        plt.savefig(save_path, bbox_inches='tight')
        print(f"Visualization saved to {save_path}")
    else:
        plt.show()

def load_model(model_path, device):
    """加载预训练模型"""
    model = ResNet18().to(device)
    checkpoint = torch.load(model_path, map_location=device)
    model.load_state_dict(checkpoint['model_state_dict'])
    model.eval()
    return model

if __name__ == "__main__":
    # 加载数据集和类别名称
    _, val_loader, _, classes = load_cifar10()
    
    # 加载预训练模型(请确保路径正确)
    model_path = "./checkpoints/best_model.pth"
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"Model not found at {model_path}. Please train the model first.")
    
    model = load_model(model_path, config.DEVICE)
    print("Model loaded successfully.")
    
    # 可视化预测结果
    visualize_predictions(
        model=model,
        val_loader=val_loader,
        classes=classes,
        num_images=8,  # 可以调整这个数字
        save_path="./visualizations/predictions.png"  # 可选: 保存可视化结果
    )